import os
import time
import cv2
import numpy as np
import sqlite3
from datetime import datetime
# 修正导入路径：从skimage.metrics导入structural_similarity
from skimage.metrics import structural_similarity as ssim
from watchdog.observers import Observer
from watchdog.events import FileSystemEventHandler


class ImageDatabase:
    """图片数据库管理类"""

    def __init__(self, db_name="flaw_detection.db"):
        self.db_name = db_name
        self.conn = sqlite3.connect(db_name)
        self._create_tables()

    def _create_tables(self):
        """创建数据库表结构"""
        cursor = self.conn.cursor()
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS reference_images (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT UNIQUE NOT NULL,
            image_data BLOB NOT NULL,
            features BLOB,
            import_time DATETIME NOT NULL
        )
        ''')
        cursor.execute('''
        CREATE TABLE IF NOT EXISTS test_images (
            id INTEGER PRIMARY KEY AUTOINCREMENT,
            filename TEXT NOT NULL,
            image_data BLOB NOT NULL,
            frame_number INTEGER,
            analysis_result TEXT,
            flaw_type TEXT,
            confidence REAL,
            analysis_time DATETIME
        )
        ''')
        self.conn.commit()

    def add_reference_image(self, filename, image_data, features=None):
        """添加基准图片到数据库"""
        try:
            cursor = self.conn.cursor()
            cursor.execute('''
            INSERT INTO reference_images (filename, image_data, features, import_time)
            VALUES (?, ?, ?, ?)
            ''', (filename, image_data, features, datetime.now()))
            self.conn.commit()
            return True
        except sqlite3.IntegrityError:
            print(f"图片 {filename} 已存在于数据库中，跳过...")
            return False

    def get_reference_images(self):
        """获取所有基准图片"""
        cursor = self.conn.cursor()
        cursor.execute("SELECT id, filename, image_data, features FROM reference_images")
        return cursor.fetchall()

    def add_test_result(self, filename, image_data, frame_number, analysis_result,
                        flaw_type=None, confidence=None):
        """添加测试结果到数据库"""
        cursor = self.conn.cursor()
        cursor.execute('''
        INSERT INTO test_images 
        (filename, image_data, frame_number, analysis_result, flaw_type, confidence, analysis_time)
        VALUES (?, ?, ?, ?, ?, ?, ?)
        ''', (filename, image_data, frame_number, analysis_result, flaw_type,
              confidence, datetime.now()))
        self.conn.commit()

    def close(self):
        """关闭数据库连接"""
        if hasattr(self, 'conn') and self.conn:
            self.conn.close()
            self.conn = None

    def __del__(self):
        """析构函数确保连接关闭"""
        self.close()


class ImageImportHandler(FileSystemEventHandler):
    """图片导入事件处理器"""

    def __init__(self, db, import_dir, timeout=30):
        self.db = db
        self.import_dir = import_dir
        self.last_import_time = time.time()
        self.timeout = timeout
        self.importing_complete = False

    def on_created(self, event):
        """新文件创建时触发"""
        if self.importing_complete:
            return
        if not event.is_directory:
            file_path = event.src_path
            if self._is_real_image(file_path):
                print(f"检测到新图片: {os.path.basename(file_path)}")
                self._import_image_safe(file_path)
                self.last_import_time = time.time()

    def check_import_complete(self):
        """检查导入是否完成"""
        return not self.importing_complete and (time.time() - self.last_import_time > self.timeout)

    def mark_as_complete(self):
        """标记导入完成"""
        self.importing_complete = True
        print("已确认参考图片导入完成")

    @staticmethod
    def _is_real_image(file_path):
        """判断是否为真实图片"""
        image_extensions = ['.jpg', '.jpeg', '.png', '.bmp', '.gif', '.tiff']
        ext = os.path.splitext(file_path)[1].lower()
        if ext not in image_extensions:
            return False

        image_magic = {
            '.jpg': b'\xff\xd8\xff',
            '.jpeg': b'\xff\xd8\xff',
            '.png': b'\x89PNG\r\n\x1a\n',
            '.bmp': b'BM',
            '.gif': (b'GIF87a', b'GIF89a'),
            '.tiff': (b'II*\x00', b'MM\x00*')
        }

        try:
            with open(file_path, 'rb') as f:
                header = f.read(8)
            if ext in ['.jpg', '.jpeg']:
                return header.startswith(image_magic[ext])
            elif ext == '.png':
                return header.startswith(image_magic[ext])
            elif ext == '.bmp':
                return header.startswith(image_magic[ext])
            elif ext == '.gif':
                return header.startswith(image_magic[ext][0]) or header.startswith(image_magic[ext][1])
            elif ext == '.tiff':
                return header.startswith(image_magic[ext][0]) or header.startswith(image_magic[ext][1])
            return False
        except Exception:
            return False

    def _import_image_safe(self, file_path):
        """安全导入图片"""
        max_attempts = 10
        attempt = 0
        image_data = None
        while attempt < max_attempts:
            try:
                with open(file_path, 'rb') as f:
                    image_data = f.read()
                break
            except PermissionError:
                attempt += 1
                time.sleep(0.5)
            except Exception as e:
                print(f"读取图片失败: {str(e)}")
                return

        if not image_data:
            print(f"图片 {os.path.basename(file_path)} 导入失败")
            return

        try:
            features = self._extract_image_features(file_path)
            filename = os.path.basename(file_path)
            self.db.add_reference_image(filename, image_data, features)
            print(f"成功导入图片: {filename}")
        except Exception as e:
            print(f"图片入库失败: {str(e)}")

    @staticmethod
    def _extract_image_features(image_path):
        """提取图片特征"""
        try:
            img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
            if img is None:
                raise ValueError("无法解析为图片")
            img = cv2.resize(img, (200, 200))
            hist = cv2.calcHist([img], [0], None, [256], [0, 256])
            hist = cv2.normalize(hist, hist).flatten()
            return hist.tobytes()
        except Exception as e:
            print(f"提取特征失败: {str(e)}")
            return None


class FlawDetector:
    """伤痕检测类"""

    def __init__(self, db):
        self.db = db
        self.reference_images = self._load_reference_features()
        if not self.reference_images:
            raise ValueError("无有效参考图片！请重新导入。")

    def _load_reference_features(self):
        """加载参考图片特征"""
        ref_images = self.db.get_reference_images()
        valid_features = []
        for img_id, filename, img_data, feat_data in ref_images:
            if not feat_data:
                continue
            try:
                feat_array = np.frombuffer(feat_data, dtype=np.float32)
                if len(feat_array) == 256:
                    valid_features.append((img_id, filename, img_data, feat_array))
            except Exception:
                continue
        print(f"已加载 {len(valid_features)} 张有效参考图片")
        return valid_features

    def process_test_images(self, test_dir):
        """处理测试图片"""
        if not os.path.exists(test_dir):
            print(f"测试目录不存在: {test_dir}")
            return []

        test_files = [os.path.join(test_dir, f) for f in os.listdir(test_dir)
                      if os.path.isfile(os.path.join(test_dir, f)) and
                      ImageImportHandler._is_real_image(os.path.join(test_dir, f))]

        if not test_files:
            print("无有效测试图片")
            return []

        results = []
        for idx, file_path in enumerate(test_files, 1):
            filename = os.path.basename(file_path)
            print(f"\n处理测试图片 {idx}/{len(test_files)}: {filename}")
            frames = self._split_into_frames(file_path)
            for frame_num, frame_data in enumerate(frames, 1):
                result = self._analyze_frame(frame_data, filename, frame_num)
                results.append(result)
        return results

    @staticmethod
    def _split_into_frames(image_path):
        """分解图片为帧"""
        try:
            img = cv2.imread(image_path)
            return [img] if img is not None else []
        except Exception as e:
            print(f"分解图片失败: {str(e)}")
            return []

    def _analyze_frame(self, frame, filename, frame_num):
        """分析单帧"""
        try:
            gray_frame = cv2.cvtColor(frame, cv2.COLOR_BGR2GRAY)
            resized_frame = cv2.resize(gray_frame, (200, 200))
            frame_hist = cv2.calcHist([resized_frame], [0], None, [256], [0, 256])
            frame_hist = cv2.normalize(frame_hist, frame_hist).flatten()

            similarities = []
            for _, ref_name, ref_data, ref_hist in self.reference_images:
                hist_sim = cv2.compareHist(frame_hist, ref_hist, cv2.HISTCMP_CORREL)
                ref_img = cv2.imdecode(np.frombuffer(ref_data, np.uint8), cv2.IMREAD_GRAYSCALE)
                ref_img = cv2.resize(ref_img, (200, 200))
                struct_sim = ssim(resized_frame, ref_img)
                similarities.append((ref_name, (hist_sim + struct_sim) / 2))

            similarities.sort(key=lambda x: x[1], reverse=True)
            best_match, best_sim = similarities[0]
            result, flaw_type, confidence = self._determine_flaw(best_sim)

            _, img_encoded = cv2.imencode('.png', frame)
            self.db.add_test_result(filename, img_encoded.tobytes(), frame_num, result, flaw_type, confidence)
            return {
                'filename': filename, 'frame': frame_num, 'result': result,
                'flaw_type': flaw_type, 'confidence': confidence, 'best_match': best_match
            }
        except Exception as e:
            return {'filename': filename, 'frame': frame_num, 'result': f"分析失败: {e}",
                    'flaw_type': None, 'confidence': 0, 'best_match': None}

    @staticmethod
    def _determine_flaw(similarity):
        """判断伤痕类型"""
        if similarity > 0.95:
            return "未检测到明显伤痕", None, 1.0
        elif similarity > 0.85:
            return "可能存在轻微划痕", "轻微划痕", 0.7
        elif similarity > 0.75:
            return "可能存在中度损伤", "中度损伤", 0.85
        else:
            return "存在明显严重损伤", "严重损伤", 1.0


def main():
    """主函数"""
    print("===== 智能探伤系统 =====")
    REF_DIR = "reference_images"
    TEST_DIR = "test_images"
    DB_NAME = "flaw_detection.db"
    TIMEOUT = 30

    os.makedirs(REF_DIR, exist_ok=True)
    os.makedirs(TEST_DIR, exist_ok=True)

    try:
        db = ImageDatabase(DB_NAME)
    except Exception as e:
        print(f"数据库初始化失败: {e}")
        return

    try:
        print(f"\n===== 导入参考图片 =====")
        print(f"请将参考图片复制到: {os.path.abspath(REF_DIR)}")
        handler = ImageImportHandler(db, REF_DIR, TIMEOUT)

        # 导入已有图片
        for f in os.listdir(REF_DIR):
            path = os.path.join(REF_DIR, f)
            if os.path.isfile(path) and ImageImportHandler._is_real_image(path):
                handler._import_image_safe(path)

        # 监控目录
        observer = Observer()
        observer.schedule(handler, REF_DIR, recursive=False)
        observer.start()

        # 等待确认
        print("等待图片导入...（按Ctrl+C强制确认）")
        while True:
            if handler.check_import_complete():
                if input(f"{TIMEOUT}秒无新图片，确认导入完成? (y/n): ").lower() in ['y', 'yes']:
                    handler.mark_as_complete()
                    break
                handler.last_import_time = time.time()
            time.sleep(1)
        observer.stop()
        observer.join()

        # 处理测试图片
        input(f"\n请将测试图片复制到: {os.path.abspath(TEST_DIR)}\n按回车开始分析...")
        detector = FlawDetector(db)
        results = detector.process_test_images(TEST_DIR)

        # 输出结果
        print("\n===== 探伤结果 =====")
        for i, res in enumerate(results, 1):
            print(f"\n测试图片 {i}: {res['filename']} (帧: {res['frame']})")
            print(f"结果: {res['result']}")
            if res['flaw_type']:
                print(f"类型: {res['flaw_type']}, 置信度: {res['confidence']:.2f}")

    except KeyboardInterrupt:
        print("\n用户中断程序")
    except Exception as e:
        print(f"程序错误: {str(e)}")
    finally:
        db.close()
        print("\n系统已关闭")


if __name__ == "__main__":
    main()